"""
@Filename       : cdk.py
@Create Time    : 2021/1/12 8:54
@Author         : Rylynn
@Description    : 

"""
import math
import os
import random

import torch as th
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from utils.dataloader import DiffuseSequenceDataset, sequence_collate_fn


class CDK(nn.Module):
    def __init__(self, config):
        super(CDK, self).__init__()
        self.n = config['embed_size']
        self.user_embed = nn.Embedding(config['user_num'], config['embed_size'])

    def forward(self, source, ui, uj):
        source_embed = self.user_embed(source)
        ui_embed = self.user_embed(ui)
        uj_embed = self.user_embed(uj)

        di = th.norm(source_embed - ui_embed, 2)
        dj = th.norm(source_embed - uj_embed, 2)
        loss = 1 - (dj - di)
        loss = th.clamp_min(loss, 0)
        return loss

    def ll(self):
        ...


def train(root_path, dataset):
    config = {
        'user_num': 4378,
        'embed_size': 64,
        'lr': 0.01,
        'epoches': 20,
    }

    train_dataloader = DataLoader(DiffuseSequenceDataset(root_path=root_path, dataset=dataset, train=True),
                                  batch_size=1,
                                  shuffle=True,
                                  collate_fn=sequence_collate_fn)

    test_dataloader = DataLoader(DiffuseSequenceDataset(root_path=root_path, dataset=dataset, train=False),
                                 batch_size=128,
                                 num_workers=4,
                                 collate_fn=sequence_collate_fn)

    model = CDK(config)
    model = model.cuda()
    optimizer = th.optim.SGD(params=model.parameters(), lr=config['lr'])
    for epoch in range(config['epoches']):
        batch_count = 0
        epoch_loss = 0
        for sequence, _, seq_length in tqdm(train_dataloader):
            sequence = sequence[0]
            seq_length = seq_length[0]

            # sample pos and neg user pairs from diffusion sequences
            source_user = sequence[0]
            pos_idx = random.randint(1, seq_length - 1)
            pos = sequence[pos_idx]
            neg = random.randint(1, config['user_num'] - 1)
            while neg in sequence[:pos_idx]:
                neg = random.randint(1, config['user_num'] - 1)
            source_user = source_user.cuda()
            pos = pos.cuda()
            neg = th.LongTensor([neg]).cuda()
            loss = model(source_user, pos, neg)

            optimizer.zero_grad()
            batch_count += 1
            epoch_loss += loss.item()

            loss.backward()
            optimizer.step()

        print('Epoches {}/{}, with average loss of {}'.format(epoch + 1, config['epoches'], epoch_loss / batch_count))
        if epoch % 3 == 0:
            ...

        ...


if __name__ == '__main__':
    train('../../data', 'memes')
